package cn.echoparrot.application;

import cn.echoparrot.domain.entity.Dataset;
import cn.echoparrot.domain.entity.Segment;

import cn.echoparrot.domain.service.DatasetService;
import cn.echoparrot.repository.neo4j.DatasetMapper;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.List;
import java.util.Optional;


@RestController
@CrossOrigin
@RequestMapping("/retrieval")
public class RetrievalController {

    private final DatasetService datasetService;
    private final DatasetMapper datasetMapper;
    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    public RetrievalController(DatasetService datasetService, DatasetMapper datasetMapper) {
        this.datasetService = datasetService;
        this.datasetMapper = datasetMapper;
    }

    @GetMapping("/{datasetId}")
    public List<Segment> retrieve(@PathVariable Long datasetId, @RequestParam String query) {
        Optional<Dataset> dataset = datasetMapper.findById(datasetId);

        return datasetService.similaritySearch(dataset.get().indexName(), 5, embeddingModel.embed(query), 0.6);
    }


}
